#!/usr/local/bin/python
# This is protocol2 - python edition.
# Perform a local alignment using TSS, cut sequences, run GSAT, run HMMTOP, and highlight alignments. Thats all.
# Requires SSearch36 From the FASTA package somewhwere in your $PATH variable.
# Created by Vamsee Reddy.

import hmmtop
import gsat
import tss,os
import pickle
import optparse
import tempfile
import hmmgap
import logging
import templates
import re
from ProjectBio import Emboss
from Bio import SeqIO
from progressbar import ProgressBar,Bar,Percentage
os.environ['BIOV_DEBUG'] = 'True'
class Compare:
	
	def __init__(self):
		# Runtime variables
		self.subject_file = None
		self.subject_name = None
		self.target_file = None
		self.target_name = None
		self.outdir = 'protocol2'
		self.assign = 3
		self.tss_shuffle = 200
		self.gsat_shuffle = 300
		self.swap = False
		self.verbose_switch = False
		self.srestrict = None
		self.trestrict = None
		
		# Internal Matricies
		self.tssdata = {}
		self.subjects_dict = {}
		self.targets_dict = {}
		self.tms = {}
		self.globaldata = []
		self.stats = {}
		
		# Storage variables
		self.tssout = '/tss.matrix'
		self.gsatout ='/gsat.matrix'
		
		# Alignment variables
		self.gapopen = 8
		self.gapextend = 2
		
		#self.init_logging()
		
	def __call__(self):
		if os.path.exists(self.outdir) is False:
			os.mkdir(self.outdir)
		self.tssearch()
		# Swap subjects & targets if needed - subject is always shorter
		if self.swap:
			(self.subject_file,self.target_file) = (self.target_file,self.subject_file)
			(self.subject_name,self.target_name) = (self.target_name,self.subject_name)
			(self.srestrict,self.trestrict) = (self.trestrict,self.srestrict)
		# Set subject & target names
		if self.subject_name is None:
			self.subject_name = self.subject_file.split('/')[-1]
		if self.target_name is None:
			self.target_name = self.target_file.split('/')[-1]
		# Write selected subjects & targets, then run HMMTOP
		subject_ids = self.tssdata.keys()
		target_ids = []
		for i in self.tssdata.values():
			target_ids.extend([x['target_symbol'] for x in i])
		target_ids = list(set(target_ids))
		subjects = SeqIO.parse(self.subject_file,'fasta')
		self.subjects_dict = SeqIO.to_dict(subjects)
		targets = SeqIO.parse(self.target_file,'fasta')
		self.targets_dict = SeqIO.to_dict(targets)
		sf = open(self.outdir+'/subjects.faa','wb+')
		tf = open(self.outdir+'/targets.faa','wb+')
		for sid in subject_ids:
			SeqIO.write(self.subjects_dict[sid],sf,'fasta')
		for tid in target_ids:
			SeqIO.write(self.targets_dict[tid],tf,'fasta')
		sf.close()
		tf.close()
		hmt = hmmtop.tools()
		hmt.add_library('subjects',self.outdir+'/subjects.faa')
		hmt.add_library('targets',self.outdir+'/targets.faa')
		hmt.scan_libraries()
		self.tms = hmt.results
		# Run GSAT on all of these
		self.run_global_alignments()
		# Sort global data
		self.globaldata.sort(key=lambda x:float(x['gsat_score']),reverse=True)
		for ii,i in enumerate(self.globaldata):
			try:
				z = int(i['gsat_score'])
				z = 0 if z<0 else z
			except:
				z = 0
			self.globaldata[ii]['gsat_score'] = z
		self.globaldata.sort(key=lambda x:int(x['gsat_score']),reverse=True)
		# Generate HTML report
		self.generate_report()
		
	def verbose(self,msg):
		if self.verbose_switch:
			print msg
			
	def init_logging(self):
		logger = logging.getLogger('PROTOCOL2')
		hdlr = logging.FileHandler('/var/tmp/PROTOCOL2.log')
		formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
		hdlr.setFormatter(formatter)
		logger.addHandler(hdlr) 
		logger.setLevel(logging.WARNING)
		self.logger = logger
		
	def tssearch(self):
		if os.path.exists(self.outdir+self.tssout):
			self.verbose('Loading existing Alignment Matrix...')
			handle = open(self.outdir+self.tssout,'r')
			try:
				(self.tssdata,self.swap) = pickle.load(handle)
				handle.close()
				return
			except ValueError:
				# Corrupt pickle - Trash it, make another.
				os.remove(handle.name)
			handle.close()
		self.verbose('Building Alignment Matrix...')
		local = tss.compare()
		local.subject = self.subject_file
		local.target = self.target_file
		local.shuffle = self.tss_shuffle
		local.max = self.assign
		local()
		# Process our data into a subject->data dict
		for result in local.results:
			[self.tssdata.setdefault(result['subject_symbol'],[]).append(result)]
			
		mymatrix = [self.tssdata,local.swap]
		self.swap = local.swap
		handle = open(self.outdir+self.tssout,'wb')
		pickle.dump(mymatrix,handle)
		handle.close()
		return
		
	def run_global_alignments(self):
		if os.path.exists(self.outdir+'/'+self.gsatout):
			handle = open(self.outdir+'/'+self.gsatout,'r')
			try:
				self.globaldata = pickle.load(handle)
				handle.close()
				return
			except ValueError:
				os.remove(handle.name)
			handle.close()
		gsatcmd = gsat.cmd()
		gsatcmd.gapopen = self.gapopen
		gsatcmd.gapextend = self.gapextend
		emboss = Emboss()
		sub_annotate = hmmgap.annotate()
		tar_annotate = hmmgap.annotate()
		sub_annotate.hmmtop = self.tms['subjects']
		tar_annotate.hmmtop = self.tms['targets']
		total = len(self.tssdata.items()) * self.assign
		i=0
		self.verbose('Performing global alignments')
		if self.verbose_switch:
			pbar = ProgressBar(widgets=[Percentage(), Bar()], maxval=total).start()
		for subject,targets in self.tssdata.items():
			subject_seq = self.subjects_dict[subject]
			for target in targets:
				mytarget = self.targets_dict[target['target_symbol']].seq[target['target_start']-1:target['target_end']]
				mysubject = subject_seq.seq[target['subject_start']-1:target['subject_end']]
				gsatcmd.asequence = str(mysubject)
				gsatcmd.bsequence = str(mytarget)
				try:
					gsatcmd()
				except:
					error = 'Subject: %s | Target: %s // \n >subject\n%s\n\n>target\n%s'%(self.subject_file,self.target_file,mysubject,mytarget)
					print error
					#self.logger.error(error)
					continue
				(subject_aln,match_aln,target_aln) = emboss.extract_alignment(gsatcmd.outfile.name)
				subject_annotation = sub_annotate(self.subjects_dict[subject],subject_aln)
				target_annotation = tar_annotate(self.targets_dict[target['target_symbol']],target_aln)
				target['gsat_score'] = gsatcmd.zscore
				(target['subject_aln'],target['match_aln'],target['target_aln']) = (subject_aln,match_aln,target_aln)
				(target['subject_annotation'],target['target_annotation']) = (subject_annotation,target_annotation)
				target['subject_tmc'] = len(self.tms['subjects'][subject]) if subject in self.tms['subjects'] else 0
				target['target_tmc'] = len(self.tms['targets'][target['target_symbol']]) if target['target_symbol'] in self.tms['targets'] else 0
				self.globaldata.append((target))
				gsatcmd.outfile.close()
				if self.verbose_switch:
					pbar.update(1+i)
					i +=1
		handle = open(self.outdir+self.gsatout,'wb')
		pickle.dump(self.globaldata,handle)
		handle.close()
		
	def calculate_overlap(self,subject,target,sd,mo=0.5):
		soverlap,toverlap = {},{}
		tms = re.compile(r'(\d+)[+_]+\+')
		sub = tms.finditer(subject)
		tar = tms.finditer(target)
		for s in sub:
			soverlap.setdefault(s.groups()[0],(s.start(),s.end()))
		for t in tar:
			toverlap.setdefault(t.groups()[0],(t.start(),t.end()))
		for stms,srange in soverlap.items():
			for ttms,trange in toverlap.items():
				# Get min len
				srange,trange=range(srange[0],srange[1]),range(trange[0],trange[1])
				minlen=[len(srange),len(trange)]
				minlen.sort()
				minlen = int(minlen[0]*mo)
				if len( set(srange)&set(trange) ) >= minlen:
					overlap = [int(stms),int(ttms)]
					overlap.sort()
					overlap = "-".join([str(i) for i in overlap])
					self.stats.setdefault(overlap,[])
					self.stats[overlap].append(sd)
			
		
	def generate_report(self):
		thandle = open(os.environ['PROTOCOL2_TEMPLATE'],'r')
		template = thandle.read()
		thandle.close()
		row_template = re.search('{ROWS}(.+?){\/ROWS}',template,re.DOTALL).groups()[0]
		data_template = re.search('{DATA}(.+?){\/DATA}',template,re.DOTALL).groups()[0]
		(myrows,mydatas,mystats) = ([],[],[])
		for result in self.globaldata:
			if self.srestrict is not None and self.trestrict is not None:
				if int(result['subject_tmc']) != self.srestrict or int(result['target_tmc']) != self.trestrict:
					continue
			myrow = row_template[:]
			mydata = data_template[:]
			result['match_aln'] = str(result['match_aln']).replace('.',' ')
			myrow = myrow.replace('{SUBJECTID}',result['subject_symbol'])
			myrow = myrow.replace('{TARGETID}',result['target_symbol'])
			myrow = myrow.replace('{SSZ}',str(round(result['z_score'],2)))
			myrow = myrow.replace('{GSATZ}',str(result['gsat_score']))
			mydata = mydata.replace('{SUBJECTID}',result['subject_symbol'])
			mydata = mydata.replace('{TARGETID}',result['target_symbol'])
			mydata = mydata.replace('{SSZ}',str(round(result['z_score'],2)))
			mydata = mydata.replace('{GSATZ}',str(result['gsat_score']))
			mydata = mydata.replace('%ATMS%',str(result['subject_annotation']))
			mydata = mydata.replace('%BTMS%',str(result['target_annotation']))
			mydata = mydata.replace('%ASEQUENCE%',str(result['subject_aln']))
			mydata = mydata.replace('%BSEQUENCE%',str(result['target_aln']))
			mydata = mydata.replace('%MATCH%',str(result['match_aln']))
			mydata = mydata.replace('%STMC%',str(result['subject_tmc']))
			mydata = mydata.replace('%TTMC%',str(result['target_tmc']))
			try:
				if int(result['gsat_score']) >= 10:
					self.calculate_overlap(result['subject_annotation'],result['target_annotation'],result['gsat_score'])
			except:
				pass
			myrows.append(myrow)
			mydatas.append(mydata)
		template = template.replace('{SUBJECTNAME}',self.subject_name)
		template = template.replace('{TARGETNAME}',self.target_name)
		rows = "".join(myrows)
		data = ''.join(mydatas)
		template = re.sub('{ROWS}.+?{\/ROWS}',rows,template,flags=re.DOTALL)
		template = re.sub('{DATA}.+?{\/DATA}',data,template,flags=re.DOTALL)
		stats = self.stats.items()
		stats.sort(key=lambda x:len(x[-1]),reverse=True)
		for key,count in stats:
			line="%s\t%i"%(key,len(count))
			mystats.append(line)
		stats = "\n".join(mystats)
		template = re.sub('%STATS%',stats,template)
		output = open(self.outdir+'/report.html','wb+')
		output.write(template)
		output.close()
		return
		


if __name__=='__main__':
	desc = "Welcome to Protocol2. This tool will allow you to rapidly locate homologs between two fasta files. \
	'Subject', 'Target', 'Outdir' Are the only mandatory options. --Subject & --Target are used to label items on\
	your actual report. EX. protocol2 -s 2.A.1.faa -t 2.A.3.faa -o mydir --subject='APC' --target='MFS' // Developed by Vamsee Reddy"
	ver = '3.0'
	opts = optparse.OptionParser(description=desc,version=ver)
	opts.add_option('-s',action='store',dest='subject',type='string',default='subject.faa',help='Path to subject file')
	opts.add_option('-t',action='store',dest='target',type='string',default='target.faa',help='Path to target file')
	opts.add_option('-o',action='store',dest='outdir',type='string',default='protocol2',help='Path to output')
	opts.add_option('--subject',action='store',dest='sname',type='string',default=None,help='Subject Name to appear on report')
	opts.add_option('--target',action='store',dest='tname',type='string',default=None,help='Target Name to appear on report')
	opts.add_option('--assign',action='store',dest='num',type='int',default=3,help='TSS Setting :: Number of targets to assign each subject')
	opts.add_option('--shuffle',action='store',dest='rand',type='int',default=300,help='Number of times to shuffle each alignment with GSAT')
	opts.add_option('--stms',action='store',dest='srestrict',type='int',default=None,help='Report will contain X TMSs subjects. TTMS must be set to work')
	opts.add_option('--ttms',action='store',dest='trestrict',type='int',default=None,help='Report will contain X TMSs targets. STMS must be set to work')
	(cli,args) = opts.parse_args()
	if os.path.exists(cli.subject) is False or os.path.exists(cli.target) is False:
		opts.print_help()
		exit()
 	protocol = Compare()
	protocol.subject_file=cli.subject
	protocol.target_file=cli.target
	protocol.subject_name=cli.sname
	protocol.target_name=cli.tname
	protocol.outdir=cli.outdir
	protocol.assign=cli.num
	protocol.gsat_shuffle=cli.rand
	protocol.verbose_switch=True
	protocol.srestrict=cli.srestrict
	protocol.trestrict=cli.trestrict
	protocol()
